from fastapi import FastAPI
from fastapi.responses import HTMLResponse
import gradio as gr
import os
from html_string import main_html,plain_html
from upload_file import *
from create_kb import *
from milvus_kb import *
from chat import get_model_response
def user(user_message, history):
    print(user_message,history)
    # 返回正确的格式：清空输入框，并添加用户消息到history
    # 对于messages格式，history应该是包含role和content的字典列表
    if history is None:
        history = []
    
    # 添加用户消息到history
    user_msg = {"role": "user", "content": user_message['text']}
    history.append(user_msg)
    
    return {'text': '','files': user_message['files']}, history

#####################################
######       gradio界面       #######
#####################################

def get_chat_block():
    with gr.Blocks(theme=gr.themes.Base(),css=".gradio_container { background-color: #f0f0f0; }") as chat:
        gr.HTML(plain_html)
        with gr.Row():     
            with gr.Column(scale=10):
                chatbot = gr.Chatbot(label="Chatbot",height=750,avatar_images=("images/user.jpeg","images/tongyi.png"),type="messages")
                with gr.Row():
                    # 
                    input_message = gr.MultimodalTextbox(label="请输入",file_types=[".xlsx",".csv",".docx",".pdf",".txt",".doc",".xls",".ppt",".pptx"],scale=7)
                    clear_btn = gr.ClearButton(chatbot,input_message,scale=1)
            # 模型与知识库参数
            with gr.Column(scale=5):
                # 合并本地和Milvus知识库列表
                def get_all_knowledge_bases():
                    local_dbs = os.listdir(DB_PATH)
                    milvus_dbs = get_milvus_collections()
                    return local_dbs + milvus_dbs
                
                knowledge_base =gr.Dropdown(choices=get_all_knowledge_bases(),label="加载知识库",interactive=True,scale=2)
                with gr.Accordion(label="召回文本段",open=False):
                    chunk_text = gr.Textbox(label="召回文本段",interactive=False,scale=5,lines=10)
                with gr.Accordion(label="股票查询功能",open=True):
                    gr.Markdown("""
                    **股票历史数据查询功能：**
                    
                    请使用明确的格式查询股票历史行情数据，例如：
                    - "查询000001.SZ从20250101到20250131的历史数据"
                    - "获取平安银行2025年1月到3月的股价数据"
                    - "查看600000.SH从20240101到20240131的历史行情"
                    
                    支持的股票代码格式：
                    - 深市：000001.SZ（平安银行）
                    - 沪市：600000.SH（浦发银行）
                    
                    日期格式：YYYYMMDD（如：20250101）
                    
                    **注意：请使用明确的日期格式，避免模糊表述**
                    """)
                with gr.Accordion(label="模型设置",open=True):
                    model =gr.Dropdown(choices=['qwen-max','qwen-plus','qwen-turbo'],label="选择模型",interactive=True,value="qwen-max",scale=2)
                    temperature = gr.Slider(maximum=2,minimum=0,interactive=True,label="温度参数",step=0.01,value=0.85,scale=2)
                    max_tokens = gr.Slider(maximum=8000,minimum=0,interactive=True,label="最大回复长度",step=50,value=1024,scale=2)
                    history_round = gr.Slider(maximum=30,minimum=1,interactive=True,label="携带上下文轮数",step=1,value=3,scale=2)
                with gr.Accordion(label="RAG参数设置",open=True):
                    chunk_cnt = gr.Slider(maximum=20,minimum=1,interactive=True,label="选择召回片段数",step=1,value=5,scale=2)
                    similarity_threshold = gr.Slider(maximum=1,minimum=0,interactive=True,label="相似度阈值",step=0.01,value=0.2,scale=2)
        input_message.submit(fn=user,inputs=[input_message,chatbot],outputs=[input_message,chatbot],queue=False).then(
            fn=get_model_response,inputs=[input_message,chatbot,model,temperature,max_tokens,history_round,knowledge_base,similarity_threshold,chunk_cnt],outputs=[chatbot,chunk_text]
            )
        chat.load(lambda: gr.update(choices=get_all_knowledge_bases()),[],knowledge_base)
        chat.load(clear_tmp)
    return chat


def get_upload_block():
    with gr.Blocks(theme=gr.themes.Base()) as upload:
        gr.HTML(plain_html)
        with gr.Tab("非结构化数据"):
            with gr.Accordion(label="新建类目",open=True):
                with gr.Column(scale=2):
                    unstructured_file = gr.Files(file_types=[".pdf",".docx",".txt",".doc",".ppt",".pptx"])
                    with gr.Row():
                        new_label = gr.Textbox(label="类目名称",placeholder="请输入类目名称",scale=5)
                        create_label_btn = gr.Button("新建类目",variant="primary",scale=1)
            with gr.Accordion(label="管理类目",open=False):
                with gr.Row():
                    data_label =gr.Dropdown(choices=os.listdir(UNSTRUCTURED_FILE_PATH),label="管理类目",interactive=True,scale=8,multiselect=True)
                    delete_label_btn = gr.Button("删除类目",variant="stop",scale=1)
        with gr.Tab("结构化数据"):
            with gr.Accordion(label="新建数据表",open=True):
                with gr.Column(scale=2):
                    structured_file = gr.Files(file_types=[".xlsx",".csv",".xls"])
                    with gr.Row():
                        new_label_1 = gr.Textbox(label="数据表名称",placeholder="请输入数据表名称",scale=5)
                        create_label_btn_1 = gr.Button("新建数据表",variant="primary",scale=1)
            with gr.Accordion(label="管理数据表",open=False):
                with gr.Row():
                    data_label_1 =gr.Dropdown(choices=os.listdir(STRUCTURED_FILE_PATH),label="管理数据表",interactive=True,scale=8,multiselect=True)
                    delete_data_table_btn = gr.Button("删除数据表",variant="stop",scale=1)
        delete_label_btn.click(delete_label,inputs=[data_label]).then(fn=update_label,outputs=[data_label])
        create_label_btn.click(fn=upload_unstructured_file,inputs=[unstructured_file,new_label]).then(fn=update_label,outputs=[data_label])
        delete_data_table_btn.click(delete_data_table,inputs=[data_label_1]).then(fn=update_datatable,outputs=[data_label_1])
        create_label_btn_1.click(fn=upload_structured_file,inputs=[structured_file,new_label_1]).then(fn=update_datatable,outputs=[data_label_1])
        upload.load(update_label,[],data_label)
        upload.load(update_datatable,[],data_label_1)
    return upload

def get_knowledge_base_block():
    with gr.Blocks(theme=gr.themes.Base()) as knowledge:
        gr.HTML(plain_html)
        
        # 存储选项选择
        with gr.Row():
            storage_type = gr.Radio(
                choices=["本地存储", "Milvus存储"],
                label="选择存储方式",
                value="本地存储",
                interactive=True
            )
        
        # 非结构化数据知识库
        with gr.Tab("非结构化数据"):
            with gr.Row():
                data_label_2 =gr.Dropdown(choices=os.listdir(UNSTRUCTURED_FILE_PATH),label="选择类目",interactive=True,scale=2,multiselect=True)
                knowledge_base_name = gr.Textbox(label="知识库名称",placeholder="请输入知识库名称",scale=2)
                create_knowledge_base_btn = gr.Button("确认创建知识库",variant="primary",scale=1)
        
        # 结构化数据知识库
        with gr.Tab("结构化数据"):
            with gr.Row():
                data_label_3 =gr.Dropdown(choices=os.listdir(STRUCTURED_FILE_PATH),label="选择数据表",interactive=True,scale=2,multiselect=True)
                knowledge_base_name_1 = gr.Textbox(label="知识库名称",placeholder="请输入知识库名称",scale=2)
                create_knowledge_base_btn_1 = gr.Button("确认创建知识库",variant="primary",scale=1)
        
        # 知识库管理
        with gr.Row():
            knowledge_base =gr.Dropdown(choices=os.listdir(DB_PATH),label="管理本地知识库",interactive=True,scale=2)
            delete_db_btn = gr.Button("删除本地知识库",variant="stop",scale=1)
        
        with gr.Row():
            milvus_knowledge_base =gr.Dropdown(choices=get_milvus_collections(),label="管理Milvus知识库",interactive=True,scale=2)
            delete_milvus_db_btn = gr.Button("删除Milvus知识库",variant="stop",scale=1)
        
        # 本地存储事件
        create_knowledge_base_btn.click(
            fn=lambda storage, name, labels: create_unstructured_db(name, labels) if storage == "本地存储" else create_unstructured_milvus_db(name, labels),
            inputs=[storage_type, knowledge_base_name, data_label_2]
        ).then(update_knowledge_base, outputs=[knowledge_base])
        
        create_knowledge_base_btn_1.click(
            fn=lambda storage, name, tables: create_structured_db(name, tables) if storage == "本地存储" else create_structured_milvus_db(name, tables),
            inputs=[storage_type, knowledge_base_name_1, data_label_3]
        ).then(update_knowledge_base, outputs=[knowledge_base])
        
        delete_db_btn.click(delete_db, inputs=[knowledge_base]).then(update_knowledge_base, outputs=[knowledge_base])
        
        # Milvus存储事件
        delete_milvus_db_btn.click(delete_milvus_db, inputs=[milvus_knowledge_base]).then(update_milvus_knowledge_base, outputs=[milvus_knowledge_base])
        
        # 页面加载事件
        knowledge.load(update_knowledge_base, [], knowledge_base)
        knowledge.load(update_milvus_knowledge_base, [], milvus_knowledge_base)
        knowledge.load(update_label, [], data_label_2)
        knowledge.load(update_datatable, [], data_label_3)
    return knowledge

app = FastAPI()
@app.get("/", response_class=HTMLResponse)
def read_main():
    html_content = main_html
    return HTMLResponse(content=html_content)


app = gr.mount_gradio_app(app, get_chat_block(), path="/chat")
app = gr.mount_gradio_app(app, get_upload_block(), path="/upload_data")
app = gr.mount_gradio_app(app, get_knowledge_base_block(), path="/create_knowledge_base")